import json
import os
import re
from typing import Dict, List
from hello_agents import Message
from mcp import ClientSession, StdioServerParameters, stdio_client
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

class BaseAgent:
    
    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description
        self.model_id = os.getenv("LLM_MODEL_ID")
        self.tools = []
        self.base_url = os.getenv("LLM_BASE_URL")
        self.api_key = os.getenv("LLM_API_KEY")
        self.timeout = float(os.getenv("LLM_TIMEOUT"))
        self.temperature = float(os.getenv("LLM_TEMPERATURE"))
        self._history: List[Dict[str, str]] = []
        self.server_params = StdioServerParameters(
            command="python",
            args=["../mcp/mcp_server.py"],
            env=None
        )
        self.client = OpenAI(
            base_url=self.base_url,
            api_key=self.api_key,
            timeout=self.timeout
            # temperature=self.temperature
        )
        # self.session: ClientSession = None
    
    async def _load_tools(self):
        """从MCP服务器加载工具列表"""
        try:
            async with stdio_client(self.server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    print(">>> 连接到MCP服务器成功")
                    # 获取工具schema
                    resp = await session.list_tools()
                    # print(resp)
                    self.tools = [{
                        "type": "function",
                        "function": {
                            "name": tool.name,
                            "description": tool.description,
                            "parameters": tool.inputSchema
                        }
                    } for tool in resp.tools]
                    print("可用工具：", [t["function"]["name"] for t in self.tools])
            

            # await session.__aexit__(None, None, None)
            # await transport.__aexit__(None, None, None)
        except Exception as e:
            print(f"❌ 加载MCP工具时出错: {str(e)}")
            raise
    
    def add_message(self, message: Message):
        """添加消息到历史记录"""
        self._history.append(message)
    
    def call_llm(self, prompt: str) -> str:
        try:
            res = self.client.chat.completions.create(
                model=self.model_id,
                messages=[
                    {"role": "system", "content": self.description},
                    {"role": "user", "content": prompt}
                ]
            )
            self._history.append({"role": "user", "content": prompt})
            self._history.append({"role": "assistant", "content": res.choices[0].message.content})
            return res.choices[0].message.content
        except Exception as e:
            raise ValueError(f"调用LLM时出错: {str(e)}")
        
    async def _call_mcp_tool(self, tool_name: str, **kwargs):
        """
        调用MCP服务器上的工具
        """
        try:
            async with stdio_client(self.server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    result = await session.call_tool(tool_name, kwargs)
                    return result.content
        except Exception as e:
            return f"调用MCP工具时出错: {str(e)}"
    
    async def call_llm_with_mcp_tools(self, **kwargs) -> str:
        """
        运行Agent并调用MCP工具
        """
        # transport = stdio_client(self.server_params) as (stdio, write):
        #     session = ClientSession(stdio, write).__aenter__()
        #     await session.initialize()
        #     print(">>> 连接到MCP服务器成功")
        # # stdio, write = await transport.__aenter__()
        # session = await ClientSession(stdio, write).__aenter__()
        # await session.initialize()
        async with stdio_client(self.server_params) as (stdio, write):
            async with ClientSession(stdio, write) as session:
                await session.initialize()
                print(">>> 连接到MCP服务器成功")
                while True:
                    print("\n请输入你的加法问题（如：5加7是多少？或'退出'）：")
                    user_input = input("> ")
                    if user_input.strip().lower() == '退出':
                        break
                    # 构造对话
                    messages = [
                        {"role": "system", "content": "你是一个加法助手，遇到加法问题请调用工具add，最后用自然语言回答用户。"},
                        {"role": "user", "content": user_input}
                    ]
                    while True:
                        print(f"🤖 {self.name} 正在处理: {user_input}")  
                        # transport = stdio_client(self.server_params)
                        # stdio, write = await transport.__aenter__()
                        # session = await ClientSession(stdio, write).__aenter__()
                        # await session.initialize()
                        response = self.client.chat.completions.create(
                            model=os.getenv("LLM_MODEL_ID"),
                            messages=messages,
                            tools=self.tools,
                            tool_choice="auto"
                        )
                        message = response.choices[0].message
                        messages.append(message)
                        if not message.tool_calls:
                            response = message.content
                            print("\nAI 回答：\n", message.content)
                            break
                        # 工具调用
                        for tool_call in message.tool_calls:
                            args = json.loads(tool_call.function.arguments)
                            result = await session.call_tool(tool_call.function.name, args)
                            messages.append({
                                "role": "tool",
                                "content": str(result),
                                "tool_call_id": tool_call.id
                            })
        # await session.__aexit__(None, None, None)
        # await transport.__aexit__(None, None, None)
        return response 
    
    def add_tool(self, tool_name: str) -> str: 
        if tool_name not in self.tools:
            self.tools.append(tool_name)
            return f"✅ {self.name} 成功添加工具: {tool_name}"
        else:
            return f"❌ {self.name} 工具 {tool_name} 已存在"
        
    # 在BaseAgent或MCPAgent类中实现的通用工具调用方法
    async def _call_mcp_tool(self, tool_name: str, **kwargs):
        """调用MCP服务器上的工具"""
        try:
            async with stdio_client(self.server_params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    result = await session.call_tool(tool_name, kwargs)
                    return result.content
        except Exception as e:
            return f"调用MCP工具时出错: {str(e)}"
        
    def toStr(self) -> str:
        return f"name: {self.name}, description: {self.description}, model_id: {self.model_id}, tools: {self.tools}, base_url: {self.base_url}, api_key: {self.api_key}, timeout: {self.timeout}, temperature: {self.temperature}"
        
import asyncio

async def main():
    agent = BaseAgent(
        name="MathAssistant",
        description="一个数学助手，能够计算表达式和生成单词"
    )
    
    await agent._load_tools()
    
    response = await agent.call_llm_with_mcp_tools()
    print(response)


if __name__ == "__main__":
    asyncio.run(main())